"""Contains a variant of the densenet model definition."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf

slim = tf.contrib.slim


def trunc_normal(stddev): return tf.truncated_normal_initializer(stddev=stddev)


def bn_act_conv_drp(current, num_outputs, kernel_size, scope='block'):
    """
    除了第一层外，每一层conv都要加上BN-Relu-Conv-Dropout操作
    :param current: 输入信息
    :param num_outputs: feature map的channel数
    :param kernel_size: 卷积核的size
    :param scope: 命名域
    :return:
    """
    current = slim.batch_norm(current, scope=scope + '_bn')
    current = tf.nn.relu(current)
    current = slim.conv2d(current, num_outputs, kernel_size, scope=scope + '_conv')
    current = slim.dropout(current, scope=scope + '_dropout')
    return current


def block(net, layers, growth, scope='block'):
    """
    每个3x3conv的layer前加一个1x1conv的layer: 减少feature map维度，降低计算量，融合每个通道的特征
    :param net: 输入信息
    :param layers: 循环的层数
    :param growth: 增长因子
    :param scope: 命名域
    :return:若输入的维度(?, 56, 56, 48)---48+layer*growth-->(?, 56, 56, 192)
    """
    for idx in range(layers):
        bottleneck = bn_act_conv_drp(net, 4 * growth, [1, 1],
                                     scope=scope + '_conv1x1' + str(idx))
        tmp = bn_act_conv_drp(bottleneck, growth, [3, 3],
                              scope=scope + '_conv3x3' + str(idx))
        net = tf.concat(axis=3, values=[net, tmp])
    return net


def transition_layer(inputs, num_outputs, scope='transition'):
    """
    过度层--每个block之间加1x1conv: 降维
    :param inputs: 输入图像信息
    :param num_outputs: feature map的channel数
    :param scope:
    :return:
    """
    network = bn_act_conv_drp(inputs, num_outputs, [1, 1], scope=scope+'_conv1x1')
    network = slim.avg_pool2d(network, [2, 2], stride=2, scope=scope+'_avgpool')
    return network


def densenet(images, num_classes=1001, is_training=False,
             dropout_keep_prob=0.8,
             scope='densenet'):
    """
    创建一个DenseNet model
    :param images: 输入信息
    :param num_classes: feature map的channel数
    :param is_training:
    :param dropout_keep_prob: 丢失率
    :param scope:
    :return: logits, end_points
    """
    growth = 24
    compression_rate = 0.5

    def reduce_dim(input_feature):
        return int(int(input_feature.shape[-1]) * compression_rate)

    end_points = {}

    with tf.variable_scope(scope, 'DenseNet', [images, num_classes]):
        with slim.arg_scope(bn_drp_scope(is_training=is_training,
                                         keep_prob=dropout_keep_prob)) as ssc:
            # (?, 224, 224, 3) --> (?, 112, 112, 48)
            network = slim.conv2d(images, 2 * growth, [7, 7], stride=2, padding='SAME', scope='Conv_0')
            # (?, 112, 112, 48) --> (?, 56, 56, 48)
            network = slim.max_pool2d(network, [3, 3], stride=2, padding='SAME', scope='pool_0')
            # (?, 56, 56, 48) --> (?, 56, 56, 192)
            network = block(network, 6, growth, scope='Block_0')
            # (?, 56, 56, 192)--> (?, 28, 28, 96)
            network = transition_layer(network, reduce_dim(network), scope='transition_0')
            # (?, 28, 28, 96) -->  (?, 28, 28, 384)
            network = block(network, 12, growth, scope='Block_1')
            # (?, 28, 28, 384) --> (?, 14, 14, 192)
            network = transition_layer(network, reduce_dim(network), scope='transition_1')
            # (?, 14, 14, 192) --> (?, 14, 14, 768)
            network = block(network, 24, growth, scope='Block_2')
            # (?, 14, 14, 768) --> (?, 7, 7, 384)
            network = transition_layer(network, reduce_dim(network), scope='transition_2')
            # (?, 7, 7, 384) --> (?, 7, 7, 768)
            network = block(network, 16, growth, scope='Block_3')
            network = slim.batch_norm(network, scope='batch_norm_relu')
            network = tf.nn.relu(network)
            # (?, 7, 7, 768) --> (?, 1, 1, 768)
            network = tf.reduce_mean(network, [1, 2], keep_dims=True, name='Global_average_pool')
            # (?, 1, 1, 768) --> (?, 1, 1, n_classes)
            network = slim.conv2d(network, num_classes, [1, 1], scope='conv_1')
            # n_classes
            logits = tf.squeeze(network, [1, 2], name='squeeze')
            # n_classes
            out = slim.softmax(logits, scope='prediction')
            end_points['logits'] = logits
            end_points['predictions'] = out

    return logits, end_points


def bn_drp_scope(is_training=True, keep_prob=0.8):
    keep_prob = keep_prob if is_training else 1
    with slim.arg_scope(
        [slim.batch_norm],
            scale=True, is_training=is_training, updates_collections=None):
        with slim.arg_scope(
            [slim.dropout],
                is_training=is_training, keep_prob=keep_prob) as bsc:
            return bsc


def densenet_arg_scope(weight_decay=0.004):
    """Defines the default densenet argument scope.

    Args:
      weight_decay: The weight decay to use for regularizing the model.

    Returns:
      An `arg_scope` to use for the inception v3 model.
    """
    with slim.arg_scope(
        [slim.conv2d],
        weights_initializer=tf.contrib.layers.variance_scaling_initializer(
            factor=2.0, mode='FAN_IN', uniform=False),
        activation_fn=None, biases_initializer=None, padding='same',
            stride=1) as sc:
        return sc


densenet.default_image_size = 224
